package org.example.flink.state;

import org.apache.flink.api.common.functions.MapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.common.typeinfo.Types;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;

import java.io.File;

public class Hello01StateOperator {
    public static void main(String[] args) throws Exception {
        //配置程序执行的参数
        Configuration configuration = new Configuration();
        configuration.setString("execution.savepoint.path",
                "D:\\javaEE\\msb\\flink_code_model\\ckpt\\64457fc1cb57c92249ef9c662f357d2e\\chk-6");

        StreamExecutionEnvironment environment = StreamExecutionEnvironment.getExecutionEnvironment(configuration);
        environment.setParallelism(2);
        environment.enableCheckpointing(5000);
        environment.getCheckpointConfig().setCheckpointStorage("file:///" + System.getProperty("user.dir")
                + File.separator + "ckpt");

        DataStreamSource<String> source = environment.socketTextStream("localhost", 19523);
        source.map(new YjxxtOperatorStateFunction()).print();

        environment.execute();
    }
}

class YjxxtOperatorStateFunction implements MapFunction<String, String>, CheckpointedFunction {
    //声明一个变量记数
    private int count;
    //创建一个状态对象
    private ListState<Integer> countListState;

    @Override
    public String map(String s) throws Exception {
        //更新计数器
        count++;
        return "[" + s.toUpperCase() + "][" + count + "]";
    }

    @Override
    public void snapshotState(FunctionSnapshotContext context) throws Exception {
        //清理一下历史数据
        countListState.clear();
        //保存数据
        countListState.add(count);
    }

    @Override
    public void initializeState(FunctionInitializationContext context) throws Exception {
        //创建对象描述器
        ListStateDescriptor<Integer> descriptor = new ListStateDescriptor<>("countListState", Types.INT);
        //创建对象
        this.countListState = context.getOperatorStateStore().getListState(descriptor);
        //判断是否恢复成功
        if (context.isRestored()) {
            System.out.println("{历史状态恢复成功}");
            countListState.get().forEach(e -> this.count = e);
        } else {
            System.out.println("{历史状态恢复失败}");
        }
    }
}